# Copyright (c) OpenMMLab. All rights reserved.
import argparse
import mmcv
import numpy as np
import os
import torch
import warnings

from mmcls.datasets import build_dataloader, build_dataset
from mmcls.models import build_classifier
from mmcls.utils import (auto_select_device,
                         setup_multi_processes,
                         wrap_non_distributed_model)
from mmcv import DictAction
from mmcv.runner import (init_dist, load_checkpoint,wrap_fp16_model)

import matplotlib.pyplot as plt
import torch
import numpy as np
from scipy.ndimage import gaussian_filter
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation




def visualize_sharpness_smooth(model, dataloader):
    _range = 1e-6
    n = 10
    # alpha_range = np.linspace(-_range, _range, n)
    # beta_range = np.linspace(-_range, _range, n)
    alpha_range = np.linspace(-0.75*_range, -0.5*_range, n)
    beta_range = np.linspace(0.3*_range, 0.5*_range, n)
    alpha_grid, beta_grid = np.meshgrid(alpha_range, beta_range)
    losses = np.zeros_like(alpha_grid)

    num_directions = 5
    direction1 = [sum(torch.randn_like(p) for _ in range(num_directions)) / num_directions for p in model.parameters()]
    direction2 = [sum(torch.randn_like(p) for _ in range(num_directions)) / num_directions for p in model.parameters()]
    direction1 = [d / (torch.norm(d) + 1e-16) for d in direction1]
    direction2 = [d / (torch.norm(d) + 1e-16) for d in direction2]

    original_params = [p.clone() for p in model.parameters()]
    
    for i, alpha in enumerate(alpha_range):
        for j, beta in enumerate(beta_range):
            with torch.no_grad():
                for p, d1, d2, original_p in zip(model.parameters(), direction1, direction2, original_params):
                    p.copy_(original_p + alpha * d1 + beta * d2)
            total_loss = 0
            # for batch in dataloader:
            batch = next(iter(dataloader))
            inputs = batch['img'].to(next(model.parameters()).device)
            targets = batch['gt_label'].to(next(model.parameters()).device)
            outputs = model(inputs)
            loss = torch.nn.CrossEntropyLoss()(outputs, targets)
            total_loss += loss.item()
            losses[i, j] = total_loss / len(dataloader)

    with torch.no_grad():
        for p, original_p in zip(model.parameters(), original_params):
            p.copy_(original_p)

    # 平滑损失
    losses = gaussian_filter(losses, sigma=1)

    fig = plt.figure(figsize=(10, 8))
    ax = fig.add_subplot(111, projection='3d')
    # ax.plot_surface(alpha_grid, beta_grid, losses, cmap='viridis', antialiased=True)
    ax.plot_surface(alpha_grid, beta_grid, losses, antialiased=True)
    ax.set_xlabel('Alpha')
    ax.set_ylabel('Beta')
    ax.set_zlabel('Loss')
    def update_view(num):
        ax.view_init(elev=-15, azim=num)
    ani = FuncAnimation(fig, update_view, frames=360, interval=50)  # 动态旋转视角
    ani.save('sharpness_rotation.gif', writer='imagemagick')  # 保存为 GIF 动画
    
    plt.tight_layout()
    plt.savefig('sharpness_visualization_smooth.png')
    print("Sharpness visualization saved as 'sharpness_visualization_smooth.png'")

def parse_args():
    parser = argparse.ArgumentParser(description='mmcls test model')
    parser.add_argument('config', help='test config file path')
    parser.add_argument('checkpoint', help='checkpoint file')
    parser.add_argument('--out', help='output result file')
    out_options = ['class_scores', 'pred_score', 'pred_label', 'pred_class']
    parser.add_argument(
        '--out-items',
        nargs='+',
        default=['all'],
        choices=out_options + ['none', 'all'],
        help='Besides metrics, what items will be included in the output '
        f'result file. You can choose some of ({", ".join(out_options)}), '
        'or use "all" to include all above, or use "none" to disable all of '
        'above. Defaults to output all.',
        metavar='')
    parser.add_argument(
        '--metrics',
        type=str,
        nargs='+',
        help='evaluation metrics, which depends on the dataset, e.g., '
        '"accuracy", "precision", "recall", "f1_score", "support" for single '
        'label dataset, and "mAP", "CP", "CR", "CF1", "OP", "OR", "OF1" for '
        'multi-label dataset')
    parser.add_argument('--show', action='store_true', help='show results')
    parser.add_argument(
        '--show-dir', help='directory where painted images will be saved')
    parser.add_argument(
        '--gpu-collect',
        action='store_true',
        help='whether to use gpu to collect results')
    parser.add_argument('--tmpdir', help='tmp dir for writing some results')
    parser.add_argument(
        '--cfg-options',
        nargs='+',
        action=DictAction,
        help='override some settings in the used config, the key-value pair '
        'in xxx=yyy format will be merged into config file. If the value to '
        'be overwritten is a list, it should be like key="[a,b]" or key=a,b '
        'It also allows nested list/tuple values, e.g. key="[(a,b),(c,d)]" '
        'Note that the quotation marks are necessary and that no white space '
        'is allowed.')
    parser.add_argument(
        '--metric-options',
        nargs='+',
        action=DictAction,
        default={},
        help='custom options for evaluation, the key-value pair in xxx=yyy '
        'format will be parsed as a dict metric_options for dataset.evaluate()'
        ' function.')
    parser.add_argument(
        '--show-options',
        nargs='+',
        action=DictAction,
        help='custom options for show_result. key-value pair in xxx=yyy.'
        'Check available options in `model.show_result`.')
    parser.add_argument(
        '--gpu-ids',
        type=int,
        nargs='+',
        help='(Deprecated, please use --gpu-id) ids of gpus to use '
        '(only applicable to non-distributed testing)')
    parser.add_argument(
        '--gpu-id',
        type=int,
        default=0,
        help='id of gpu to use '
        '(only applicable to non-distributed testing)')
    parser.add_argument(
        '--launcher',
        choices=['none', 'pytorch', 'slurm', 'mpi'],
        default='none',
        help='job launcher')
    parser.add_argument('--local_rank', type=int, default=0)
    parser.add_argument('--device', help='device used for testing')
    parser.add_argument(
        '--output-prediction',
        help='where to save prediction in csv file',
        default=False)
    parser.add_argument(
        '--visualize_sharpness',
        help='where to save prediction in csv file',
        default=True)
    
    args = parser.parse_args()
    if 'LOCAL_RANK' not in os.environ:
        os.environ['LOCAL_RANK'] = str(args.local_rank)

    assert args.metrics or args.out, \
        'Please specify at least one of output path and evaluation metrics.'

    return args


def main():
    args = parse_args()

    cfg = mmcv.Config.fromfile(args.config)
    if args.cfg_options is not None:
        cfg.merge_from_dict(args.cfg_options)
    cfg.output_prediction = args.output_prediction
    # set multi-process settings
    setup_multi_processes(cfg)

    # set cudnn_benchmark
    if cfg.get('cudnn_benchmark', False):
        torch.backends.cudnn.benchmark = True
    cfg.model.pretrained = None

    if args.gpu_ids is not None:
        cfg.gpu_ids = args.gpu_ids[0:1]
        warnings.warn('`--gpu-ids` is deprecated, please use `--gpu-id`. '
                      'Because we only support single GPU mode in '
                      'non-distributed testing. Use the first GPU '
                      'in `gpu_ids` now.')
    else:
        cfg.gpu_ids = [args.gpu_id]
    cfg.device = args.device or auto_select_device()

    # init distributed env first, since logger depends on the dist info.
    if args.launcher == 'none':
        distributed = False
    else:
        distributed = True
        init_dist(args.launcher, **cfg.dist_params)
    cfg.data.test.pipeline = cfg.data.train.pipeline
    dataset = build_dataset(cfg.data.test)

    # build the dataloader
    # The default loader config
    loader_cfg = dict(
        # cfg.gpus will be ignored if distributed
        num_gpus=1 if cfg.device == 'ipu' else len(cfg.gpu_ids),
        dist=distributed,
        round_up=True,
    )
    # The overall dataloader settings
    loader_cfg.update({
        k: v
        for k, v in cfg.data.items() if k not in [
            'train', 'val', 'test', 'train_dataloader', 'val_dataloader',
            'test_dataloader'
        ]
    })
    test_loader_cfg = {
        **loader_cfg,
        'shuffle': False,  # Not shuffle by default
        'sampler_cfg': None,  # Not use sampler by default
        **cfg.data.get('test_dataloader', {}),
    }
    # the extra round_up data will be removed during gpu/cpu collect
    data_loader = build_dataloader(dataset, **test_loader_cfg)

    # build the model and load checkpoint
    model = build_classifier(cfg.model)
    fp16_cfg = cfg.get('fp16', None)
    if fp16_cfg is not None:
        wrap_fp16_model(model)
    checkpoint = load_checkpoint(model, args.checkpoint, map_location='cpu')

    if 'CLASSES' in checkpoint.get('meta', {}):
        CLASSES = checkpoint['meta']['CLASSES']
    else:
        from mmcls.datasets import ImageNet
        warnings.simplefilter('once')
        warnings.warn('Class names are not saved in the checkpoint\'s '
                      'meta data, use imagenet by default.')
        CLASSES = ImageNet.CLASSES


    model = wrap_non_distributed_model(model, device=cfg.device, device_ids=cfg.gpu_ids)
    model.CLASSES = CLASSES




    if args.visualize_sharpness:
        # 确保单 GPU 模式
        if distributed:
            raise ValueError("Sharpness visualization only supports single GPU mode.")
        visualize_sharpness_smooth(model, data_loader)


if __name__ == '__main__':
    main()
